#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
北证50指数成分股数据查询工具
用于分析成分股变化历史和统计信息
"""

import sqlite3
import pandas as pd
import json
from datetime import datetime
import matplotlib.pyplot as plt
import seaborn as sns
from typing import List, Dict, Optional

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class BSE50ConstituentsAnalyzer:
    """北证50成分股数据分析器"""

    def __init__(self, db_path: str = "bse_data.db"):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)

    def __del__(self):
        if hasattr(self, 'conn'):
            self.conn.close()

    def get_all_changes(self) -> pd.DataFrame:
        """获取所有成分股变化记录"""
        query = """
        SELECT date, constituents_count, added_stocks, removed_stocks, change_type, created_at
        FROM bse50_constituents_changes
        ORDER BY date
        """
        return pd.read_sql_query(query, self.conn)

    def get_constituents_on_date(self, date: str) -> List[str]:
        """获取指定日期的成分股列表"""
        query = """
        SELECT stock_code FROM bse50_constituents_snapshot
        WHERE date = ? AND is_active = 1
        ORDER BY stock_code
        """
        cursor = self.conn.cursor()
        cursor.execute(query, (date,))
        return [row[0] for row in cursor.fetchall()]

    def get_stock_history(self, stock_code: str) -> pd.DataFrame:
        """获取指定股票的成分股历史"""
        query = """
        SELECT date, is_active
        FROM bse50_constituents_snapshot
        WHERE stock_code = ?
        ORDER BY date
        """
        return pd.read_sql_query(query, self.conn, params=(stock_code,))

    def get_change_summary(self) -> Dict:
        """获取变化摘要统计"""
        changes_df = self.get_all_changes()

        if changes_df.empty:
            return {}

        # 解析JSON字段
        changes_df['added_count'] = changes_df['added_stocks'].apply(
            lambda x: len(json.loads(x)) if pd.notna(x) else 0
        )
        changes_df['removed_count'] = changes_df['removed_stocks'].apply(
            lambda x: len(json.loads(x)) if pd.notna(x) else 0
        )
        changes_df['total_changes'] = changes_df['added_count'] + changes_df['removed_count']

        # 统计信息
        summary = {
            'total_records': len(changes_df),
            'date_range': {
                'start': changes_df['date'].min(),
                'end': changes_df['date'].max()
            },
            'total_changes': changes_df['total_changes'].sum(),
            'total_added': changes_df['added_count'].sum(),
            'total_removed': changes_df['removed_count'].sum(),
            'change_types': changes_df['change_type'].value_counts().to_dict(),
            'major_changes': changes_df[changes_df['total_changes'] > 10].shape[0]
        }

        return summary

    def get_most_frequent_constituents(self, limit: int = 20) -> pd.DataFrame:
        """获取出现频率最高的成分股"""
        query = """
        SELECT stock_code, COUNT(*) as appearance_count
        FROM bse50_constituents_snapshot
        WHERE is_active = 1
        GROUP BY stock_code
        ORDER BY appearance_count DESC
        LIMIT ?
        """
        return pd.read_sql_query(query, self.conn, params=(limit,))

    def get_current_constituents(self) -> List[str]:
        """获取当前成分股列表"""
        # 获取最新日期
        query = """
        SELECT MAX(date) FROM bse50_constituents_changes
        """
        cursor = self.conn.cursor()
        cursor.execute(query)
        latest_date = cursor.fetchone()[0]

        if latest_date:
            return self.get_constituents_on_date(latest_date)
        return []

    def print_summary(self):
        """打印分析摘要"""
        summary = self.get_change_summary()
        current_constituents = self.get_current_constituents()

        print("\n" + "="*60)
        print("北证50指数成分股分析摘要")
        print("="*60)

        if not summary:
            print("没有数据")
            return

        print(f"\n📊 基本统计:")
        print(f"  数据记录数: {summary['total_records']}")
        print(f"  时间范围: {summary['date_range']['start']} 至 {summary['date_range']['end']}")
        print(f"  总变化次数: {summary['total_changes']}")
        print(f"  新增股票总数: {summary['total_added']}")
        print(f"  移除股票总数: {summary['total_removed']}")
        print(f"  当前成分股数量: {len(current_constituents)}")

        print(f"\n🔄 变化类型分布:")
        for change_type, count in summary['change_types'].items():
            type_name = {
                'initial': '初始数据',
                'quarterly': '季度调整',
                'irregular': '非规则调整',
                'irregular_before': '调整前',
                'irregular_after': '调整后',
                'major_quarterly': '重大季度调整'
            }.get(change_type, change_type)
            print(f"  {type_name}: {count}")

        if summary['major_changes'] > 0:
            print(f"\n⚠️  大规模调整: {summary['major_changes']} 次 (单次变化>10只股票)")

        print(f"\n🏆 当前成分股 (前10只):")
        for i, stock in enumerate(current_constituents[:10], 1):
            print(f"  {i:2d}. {stock}")

    def plot_changes_over_time(self):
        """绘制成分股变化时间序列图"""
        changes_df = self.get_all_changes()

        if changes_df.empty:
            print("没有数据可以绘制")
            return

        # 解析JSON字段
        changes_df['added_count'] = changes_df['added_stocks'].apply(
            lambda x: len(json.loads(x)) if pd.notna(x) else 0
        )
        changes_df['removed_count'] = changes_df['removed_stocks'].apply(
            lambda x: len(json.loads(x)) if pd.notna(x) else 0
        )
        changes_df['total_changes'] = changes_df['added_count'] + changes_df['removed_count']

        # 转换日期格式
        changes_df['date'] = pd.to_datetime(changes_df['date'])

        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle('北证50指数成分股变化分析', fontsize=16)

        # 1. 总变化数量时间序列
        axes[0, 0].plot(changes_df['date'], changes_df['total_changes'], 'b-o', markersize=4)
        axes[0, 0].set_title('成分股变化数量时间序列')
        axes[0, 0].set_ylabel('变化数量')
        axes[0, 0].grid(True, alpha=0.3)

        # 2. 新增vs移除股票数量
        axes[0, 1].plot(changes_df['date'], changes_df['added_count'], 'g-o', label='新增', markersize=4)
        axes[0, 1].plot(changes_df['date'], changes_df['removed_count'], 'r-o', label='移除', markersize=4)
        axes[0, 1].set_title('新增vs移除股票数量')
        axes[0, 1].set_ylabel('股票数量')
        axes[0, 1].legend()
        axes[0, 1].grid(True, alpha=0.3)

        # 3. 变化类型分布
        change_types = changes_df['change_type'].value_counts()
        axes[1, 0].pie(change_types.values, labels=change_types.index, autopct='%1.1f%%')
        axes[1, 0].set_title('变化类型分布')

        # 4. 成分股数量稳定性
        axes[1, 1].plot(changes_df['date'], changes_df['constituents_count'], 'purple', linewidth=2)
        axes[1, 1].axhline(y=50, color='red', linestyle='--', alpha=0.7, label='标准数量(50)')
        axes[1, 1].set_title('成分股数量稳定性')
        axes[1, 1].set_ylabel('成分股数量')
        axes[1, 1].legend()
        axes[1, 1].grid(True, alpha=0.3)

        plt.tight_layout()
        plt.savefig('bse50_constituents_analysis.png', dpi=300, bbox_inches='tight')
        plt.show()

        print("图表已保存: bse50_constituents_analysis.png")

    def export_to_csv(self, filename: str = "bse50_constituents_export.csv"):
        """导出数据到CSV文件"""
        # 获取所有数据
        changes_df = self.get_all_changes()
        frequent_stocks = self.get_most_frequent_constituents()

        # 保存变化记录
        changes_df.to_csv(f"{filename}_changes.csv", index=False, encoding='utf-8')

        # 保存频率统计
        frequent_stocks.to_csv(f"{filename}_frequency.csv", index=False, encoding='utf-8')

        print(f"数据已导出:")
        print(f"  变化记录: {filename}_changes.csv")
        print(f"  频率统计: {filename}_frequency.csv")


def main():
    """主函数"""
    try:
        analyzer = BSE50ConstituentsAnalyzer()

        # 打印摘要
        analyzer.print_summary()

        # 绘制图表
        analyzer.plot_changes_over_time()

        # 导出数据
        analyzer.export_to_csv()

    except Exception as e:
        print(f"分析失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()